# -*- coding:utf-8 -*-
# @FileName :llm_chat.py
# @Created  :2025/1/19
# @Author   :天空之城
import time

import ollama


class OllamaChatHandler:
    def __init__(self, model_name="qwen2.5"):
        self.model_name = model_name
        self.messages = []  # 初始化对话历史列表

    def chat(self, user_input, stream=False):
        # 添加用户输入到对话历史中
        self.messages.append({"role": "user", "content": user_input})

        # 调用 Ollama 模型进行对话
        back = ollama.chat(model=self.model_name, messages=self.messages, stream=stream)

        # 获取模型的回复
        model_reply = back.message.content

        # 将模型的回复也添加到对话历史中，确保对话连贯
        self.messages.append({"role": "assistant", "content": model_reply})

        return model_reply

    def stream_chat(self, user_input):
        # 添加用户输入到对话历史中
        self.messages.append({"role": "user", "content": user_input})

        # 调用 Ollama 模型进行流式对话
        response_generator = ollama.chat(model=self.model_name, messages=self.messages, stream=True)

        # 实时输出模型的回复
        model_reply = ""  # 用于存储完整的模型回复
        for chunk in response_generator:
            if chunk.message and chunk.message.content:
                print(chunk.message.content, end='', flush=True)
                time.sleep(0.1)
                model_reply += chunk.message.content  # 累加回复内容

        # 将完整的模型回复添加到对话历史中
        self.messages.append({"role": "assistant", "content": model_reply})
        print()  # 换行


if __name__ == '__main__':
    ollama_chat_handler = OllamaChatHandler()

    # 进行第一次流式对话
    user_input1 = ("我想你帮我看一些问题，找出问题原因，并帮我解决，"
                   "我有日志、云图、调用链，以及你帮我看一下哪些工具可以用的，"
                   "如果工具不能用，可以用结合你自己的分析给出一个结论,明白了吗，明白了我发错误问题给你，你回复明白了即可")
    print("用户:", user_input1)
    ollama_chat_handler.stream_chat(user_input1)

    error_log = """
    """
    # 进行第二次流式对话
    user_input2 = error_log
    print("用户:", user_input2)
    ollama_chat_handler.stream_chat(user_input2)
